const { train } = require("@tensorflow/tfjs");
import * as tf from "@tensorflow/tfjs";
import { draw } from "./components/draw";

const host = "http://localhost:5000";
const train_path = `${host}/ocrdataset/train_annotation.txt`;
const target_height = 32;
const maxlabellength = 20;

function drawImg(data) {
  draw(document.body, data, "alpha");
}
export async function loadDataset() {
  let result = "";
  const content = await (await fetch(train_path)).text();
  const dataset = await Promise.all(generator(parseAnnotation(content)));
  const inputs = stackInput(
    dataset.reduce(
      (acc, item) => {
        acc.the_input.push(item.input);
        const label_length = item.labels.length;
        acc.label_length.push(label_length);
        acc.input_length.push(item.input_length);
        acc.the_labels.push(
          tf.pad(
            tf.tensor(item.labels, [label_length], "int32"),
            [[0, maxlabellength - label_length]],
            10000
          )
        );

        return acc;
      },
      { the_input: [], the_labels: [], input_length: [], label_length: [] }
    )
  );

  const outputs = { ctc: tf.zeros([dataset.length]) };
  return [inputs, outputs];
  // const reader = res.body.getReader();
  // let charsReceived = 0;
  // const {done, value} = await reader.read();
  // if (done) {
  //   console.log("Stream complete");
  //   return;
  // }
  // charsReceived += value.length;
  // const chunk = value;
  // result += chunk;
  // console.log(value);
}

function parseAnnotation(content) {
  return content
    .split("\n")
    .map((line) => {
      if (line.trim()) {
        const array = line.split(",");
        console.assert(array.length === 4, "annotation file format error");
        const codes = array[3].split(/\s+/).map((item) => +item);
        return {
          size: array.slice(0, 2),
          input: array[2],
          labels: codes,
        };
      }
    })
    .filter((item) => !!item);
}

function generator(annotation) {
  return annotation.map(async (anno) => {
    let imgData = await loadImg(`${host}/${anno.input}`);
    // const [h, w] = anno.size;
    // const scale = h / target_height;
    // const width = Math.floor(w / scale);
    const width = 280;
    imgData = tf.image.resizeBilinear(imgData, [target_height, width]);
    imgData = tf.mean(imgData, 2).expandDims(-1);
    drawImg(imgData);
    // todo: float32
    imgData = tf.sub(tf.div(imgData, 255), 0.5);
    return {
      input: imgData,
      labels: anno.labels,
      input_length: Math.floor(width / 8),
    };
  });
}

function stackInput(input) {
  const result = {};
  ["the_input", "the_labels", "input_length", "label_length"].forEach((key) => {
    result[key] = tf.stack(input[key]);
  });
  return result;
}

export function loadImg(url) {
  let resolve, reject;
  let promise = new Promise((reso, rej) => {
    resolve = reso;
    reject = rej;
  });
  const img = new Image();
  img.crossOrigin = "anonymous";
  img.onload = () => {
    resolve(tf.browser.fromPixels(img));
  };
  img.src = url;
  return promise;
}
